﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Roslyn.Compilers.CSharp;
using Roslyn.Compilers.Common;
using RoslynHelper;

namespace WeakEvent
{
    internal class WeakEventRewriter : SyntaxRewriter
    {
        private ISemanticModel _semanticModel;
        private SyntaxTokenList _eventModifiers;

        public WeakEventRewriter(ISemanticModel semanticModel)
        {
            _semanticModel = semanticModel;
        }

        protected override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax node)
        {
            if (new WeakEventClassVisitor().Visit(node))
            {
                string weakEventManager =
@"
public class WeakEventManager
{
    public static System.IDisposable HandleWeakly<TTarget>(
        TTarget target,
        System.Action<TTarget> weakCallback,
        System.Action<Action> addHandler,
        System.Action<Action> removeHandler)
        where TTarget : class
    {
        return new WeakEventToken<TTarget>(target, addHandler, removeHandler, weakCallback);

    }

    private class WeakEventToken<TTarget> : IDisposable
        where TTarget : class
    {
        System.WeakReference _target;
        System.Action _delegate;
        System.Action<TTarget> _callback;
        System.Action<Action> _removeHandler;
        public WeakEventToken(
            TTarget target,
            System.Action<Action> addHandler,
            System.Action<Action> removeHandler,
            System.Action<TTarget> weakCallback)
        {
            _target = new System.WeakReference(target);
            _callback =  weakCallback;
            _delegate = CallCallback;
            _removeHandler = removeHandler;
            addHandler(_delegate);
        }

        private void CallCallback()
        {
            var target = (TTarget)_target.Target;
            if (target == null)
            {
                Dispose();
                return;
            }
            _callback(target);
        }

        public void Dispose()
        {
            _removeHandler(_delegate);
            _target = null;
            _delegate = null;
            _callback = null;
            _removeHandler = null;
        }
    }
}
";
                return base.VisitClassDeclaration(node.AddMember((ClassDeclarationSyntax)Syntax.ParseCompilationUnit(weakEventManager).ChildNodes().First()));
            }
            return base.VisitClassDeclaration(node);
        }

        protected override SyntaxNode VisitBlock(BlockSyntax node)
        {
            BinaryExpressionSyntax assignment; 
            ExpressionStatementSyntax expressionStatementSyntax;
            if (WeakEventClassVisitor.IsWeakHandler(node, out assignment, out expressionStatementSyntax))
            {
                MemberAccessExpressionSyntax memberAccessExpression = assignment.Left as MemberAccessExpressionSyntax;
                if (memberAccessExpression != null)
                {
                    var classNode = node.Parent;
                    while (!(classNode is ClassDeclarationSyntax))
                        classNode = classNode.Parent;
                    return node.DefineStatements(
                        Syntax.InvocationExpression(
                            Syntax.MemberAccessExpression(
                                SyntaxKind.MemberAccessExpression,
                                Syntax.IdentifierName("WeakEventManager"),
                                name: Syntax.GenericName(
                                    Syntax.Identifier("HandleWeakly"),
                                    Syntax.TypeArgumentList(
                                    arguments: SeparatedList.Create<TypeSyntax>(
                                        Syntax.IdentifierName(((ClassDeclarationSyntax)classNode).Identifier.GetText()))))),
                        argumentList: Syntax.ArgumentList(
                            arguments: SeparatedList.Create<ArgumentSyntax>(
                                Syntax.Argument(
                                    expression: Syntax.ThisExpression()),
                                Syntax.Argument(
                                    expression: Syntax.SimpleLambdaExpression(
                                        Syntax.Parameter(
                                            identifier: Syntax.Identifier("target")),
                                        body: Syntax.InvocationExpression(
                                            Syntax.MemberAccessExpression(
                                                SyntaxKind.MemberAccessExpression,
                                                Syntax.IdentifierName("target"),
                                                name: (IdentifierNameSyntax)expressionStatementSyntax.Expression),
                                            Syntax.ArgumentList()))),
                                Syntax.Argument(
                                    expression: Syntax.SimpleLambdaExpression(
                                        Syntax.Parameter(
                                            identifier: Syntax.Identifier("h")),
                                        body: Syntax.BinaryExpression(
                                            SyntaxKind.AddAssignExpression,
                                            assignment.Left,
                                            right: Syntax.IdentifierName("h")))),
                                Syntax.Argument(
                                    expression: Syntax.SimpleLambdaExpression(
                                        Syntax.Parameter(
                                            identifier: Syntax.Identifier("h")),
                                        body: Syntax.BinaryExpression(
                                            SyntaxKind.SubtractAssignExpression,
                                            assignment.Left, 
                                            right: Syntax.IdentifierName("h"))))))));
                }
            }
            return base.VisitBlock(node);
        }

        protected override SyntaxNode VisitEventFieldDeclaration(EventFieldDeclarationSyntax node)
        {
            if (_eventModifiers != default(SyntaxTokenList))
            {
                var value = node.DefineModifiers(_eventModifiers);
                _eventModifiers = default(SyntaxTokenList);
                return value;
            }
            return base.VisitEventFieldDeclaration(node);
        }
    }
}
